#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：Employment_class 
@File    ：12_tftp_download_uploading.py
@IDE     ：PyCharm 
@Author  ：XiaoDengYa
@Date    ：2021/6/23/023 16:03 
"""


import sys
import os
import struct
from socket import *

class TftpDownloadUploading(object):
    def __init__(self):
        # if len(sys.argv) != 3:
        #     print("-" * 30)
        #     print("请输入正确格式：python xxx.py ip")
        #     print("例如：python xxx.py 127.0.0.1")
        #     print("-" * 30)
        # else:
        self.ip = input("请输入ip：")
        self.select = int(input("请选择下载或上传 0 or 1："))
        self.filename = input("请输入需要操作的文件名称：")
        self.udpSocket = socket(AF_INET,SOCK_DGRAM)

    def download_file(self):
        requestFile = struct.pack("!H%dsb5sb" % len(self.filename), 1, self.filename.encode("utf-8"), 0, "octet".encode("utf-8"), 0)
        self.udpSocket.sendto(requestFile, (self.ip,69))
        flag = True
        num = 0
        with open(self.filename,"ab+") as f:
            while True:
                responseData = self.udpSocket.recvfrom(1024)
                recvData,serverInfo = responseData
                opNum = struct.unpack("!H",recvData[:2])
                packetNum = struct.unpack("!H", recvData[2:4])
                if opNum[0] == 3:
                    num += 1
                    if num == 65536:
                        num = 0
                    if num == packetNum[0]:
                        f.write(recvData[4:])
                        num = packetNum[0]

                    ackData = struct.pack("!HH", 4, packetNum[0])
                    self.udpSocket.sendto(ackData, serverInfo)
                elif opNum[0] == 5:
                    print("error，没这个文件。。。")
                    flag = False

                if len(recvData) < 516:
                    break
        if not flag:
            os.unlink(self.filename)


    def uploading_file(self):
        requestFile = struct.pack("!H%dsb5sb" % len(self.filename), 2, self.filename.encode("utf-8"), 0,
                                  "octet".encode("utf-8"), 0)
        self.udpSocket.sendto(requestFile, (self.ip, 69))

        current_number = 1
        num = 0
        while True:
            data, addr = self.udpSocket.recvfrom(1024)
            code, block_number = struct.unpack("!HH", data[:4])
            f = open(self.filename, "rb")
            if code == 4:
                if block_number == 0:
                    f = open(self.filename, "rb")

                if block_number + 1 == current_number:
                    file_data = f.read(512)
                    send_data = struct.pack("!HH", 3, current_number) + file_data
                    self.udpSocket.sendto(send_data, addr)
                    print("第{}次传输数据！".format(current_number))
                    block_number += 1
                    current_number += 1

                if len(send_data) < 516:
                    f.close()
                    print("文件上传完成！")
                    break
            f.close()

    def main(self):
        if self.select == 0:
            print("开始进行文件下载")
            self.download_file()
        elif self.select == 1:
            print("开始进行文件上传")
            self.uploading_file()
        else:
            print("选择有误！只能选择上传或下载！")
            exit()

if __name__ == '__main__':
    TftpDownloadUploading().main()
